{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MPPI (Model Predictive Path-Integral) Control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from typing import Tuple\n",
    "from matplotlib import patches\n",
    "from matplotlib.animation import ArtistAnimation\n",
    "from IPython import display"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Control Target : CartPole\n",
    "\n",
    "<img src=\"../media/cartpole.png\" width=\"500px\" />"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CartPole():\n",
    "    def __init__(\n",
    "            self,\n",
    "            mass_of_cart: float = 1.0,\n",
    "            mass_of_pole: float = 0.01,\n",
    "            length_of_pole: float = 2.0,\n",
    "            max_force_abs: float = 100.0,\n",
    "            delta_t: float = 0.02,\n",
    "            visualize: bool = True,\n",
    "        ) -> None:\n",
    "        \"\"\"initialize cartpole environment\n",
    "        state variables:\n",
    "            x: horizontal position of the cart\n",
    "            theta: angle of the pole (positive in the counter-clockwise direction)\n",
    "            x_dot: horizontal velocity of the cart\n",
    "            theta_dot: angular velocity of the pole\n",
    "        control input:\n",
    "            force: force applied to the cart (positive in the right direction)\n",
    "        Note: dynamics of the cartpole is given by OpenAI gym implementation; https://github.com/openai/gym/blob/master/gym/envs/classic_control/cartpole.py\n",
    "        \"\"\"\n",
    "        # physical parameters\n",
    "        self.g = 9.81\n",
    "        self.mass_of_cart = mass_of_cart\n",
    "        self.mass_of_pole = mass_of_pole\n",
    "        self.length_of_pole = length_of_pole\n",
    "        self.max_force_abs = max_force_abs\n",
    "        self.delta_t = delta_t\n",
    "\n",
    "        # visualization settings\n",
    "        self.cart_w = 1.8\n",
    "        self.cart_h = 1.0\n",
    "        self.max_length_of_force_arrow = 4.0\n",
    "        self.view_x_lim_min, self.view_x_lim_max = -6.0, 6.0\n",
    "        self.view_y_lim_min, self.view_y_lim_max = -6.0, 6.0\n",
    "\n",
    "        # reset environment\n",
    "        self.visualize_flag = visualize\n",
    "        self.reset()\n",
    "\n",
    "    def reset(\n",
    "            self, \n",
    "            init_state: np.ndarray = np.array([0.0, np.pi, 0.0, 0.0]), # [x, theta, x_dot, theta_dot]\n",
    "        ) -> None:\n",
    "        \"\"\"reset environment to initial state\"\"\"\n",
    "\n",
    "        # reset state variables\n",
    "        self.state = init_state\n",
    "\n",
    "        # clear animation frames\n",
    "        self.frames = []\n",
    "\n",
    "        if self.visualize_flag:\n",
    "            # prepare figure\n",
    "            self.fig, self.ax = plt.subplots(1, 1, figsize=(9,9))\n",
    "\n",
    "            # graph layout settings\n",
    "            self.ax.set_xlim(self.view_x_lim_min, self.view_x_lim_max)\n",
    "            self.ax.set_ylim(self.view_y_lim_min, self.view_y_lim_max)\n",
    "            self.ax.tick_params(labelbottom=False, labelleft=False, labelright=False, labeltop=False)\n",
    "            self.ax.tick_params(bottom=False, left=False, right=False, top=False)\n",
    "            self.ax.set_aspect('equal')\n",
    "\n",
    "    def update(self, u: np.ndarray, delta_t: float = 0.0, append_frame=True) -> None:\n",
    "        \"\"\"update state variables\"\"\"\n",
    "        # keep previous states\n",
    "        x, theta, x_dot, theta_dot= self.state\n",
    "\n",
    "        # prepare params\n",
    "        g = self.g\n",
    "        M = self.mass_of_cart\n",
    "        m = self.mass_of_pole\n",
    "        l = self.length_of_pole\n",
    "        dt = self.delta_t if delta_t == 0.0 else delta_t\n",
    "\n",
    "        # limit input force\n",
    "        force = np.clip(u, -self.max_force_abs, self.max_force_abs)[0]\n",
    "\n",
    "        # get acc. values\n",
    "        temp = (\n",
    "            force + (m*l) * theta_dot**2 * np.sin(theta)\n",
    "        ) / (M + m)\n",
    "        new_theta_ddot = (g * np.sin(theta) - np.cos(theta) * temp) / (\n",
    "            l * (4.0 / 3.0 - m * np.cos(theta)**2 / (M + m))\n",
    "        )\n",
    "        new_x_ddot = temp - (m*l)  * new_theta_ddot * np.cos(theta) / (M+m)\n",
    "\n",
    "        # update pos. values\n",
    "        new_x = x + x_dot * dt\n",
    "        new_theta = theta + theta_dot * dt\n",
    "        new_theta = ((new_theta + np.pi) % (2 * np.pi)) - np.pi # normalize theta to [-pi, pi]\n",
    "\n",
    "        # update vel. values\n",
    "        new_theta_dot = theta_dot + new_theta_ddot * dt\n",
    "        new_x_dot = x_dot + new_x_ddot * dt\n",
    "\n",
    "        # update state variables\n",
    "        self.state = np.array([new_x, new_theta, new_x_dot, new_theta_dot])\n",
    "\n",
    "        # record frame\n",
    "        if append_frame:\n",
    "            self.append_frame(force)\n",
    "\n",
    "    def get_state(self) -> np.ndarray:\n",
    "        \"\"\"return state variables\"\"\"\n",
    "        return self.state.copy()\n",
    "\n",
    "    def append_frame(self, force) -> list:\n",
    "        \"\"\"draw a frame of the animation.\"\"\"\n",
    "        # draw the cartpole\n",
    "        x, theta, x_dot, theta_dot = self.state\n",
    "        origin_x, origin_y = x, 0.0\n",
    "        w = 0.35 # width of the pole\n",
    "        l = self.length_of_pole # length of the pole\n",
    "        e = -0.2 # param of the pole shape\n",
    "        d = 0.05 # param of the pole shape\n",
    "        pole_shape_x = [e, e, e+d, l-d, l, l, l-d, e+d, e, e]\n",
    "        pole_shape_y = [0.0, 0.5*w-d, 0.5*w, 0.5*w, 0.5*w-d, -0.5*w+d, -0.5*w, -0.5*w, -0.5*w+d, 0.0]\n",
    "        rotated_pole_shape_x, rotated_pole_shape_y = self._affine_transform(pole_shape_x, pole_shape_y, theta+0.5*np.pi, [origin_x, origin_y])\n",
    "        frame = self.ax.plot(rotated_pole_shape_x, rotated_pole_shape_y, color='black', linewidth=2.0, zorder=3)\n",
    "        frame += self.ax.fill(rotated_pole_shape_x, rotated_pole_shape_y, color='white', zorder=2)\n",
    "\n",
    "        # draw the cart and the horizontal line\n",
    "        cart_x, cart_y = x, 0.0\n",
    "        frame += [self.ax.add_artist(patches.Rectangle(xy=(cart_x-self.cart_w/2.0, cart_y-self.cart_h/2.0), width=self.cart_w, height=self.cart_h, ec=\"black\", linewidth=2.0, fc=\"white\", fill=True, zorder=1))]\n",
    "        frame += [self.ax.hlines(0,self.view_x_lim_min,self.view_x_lim_max,colors=\"gray\", zorder=0)]\n",
    "\n",
    "        # draw the joint circle\n",
    "        joint = patches.Circle([origin_x, origin_y], radius=abs(e)/2.8, fc='white', ec='black', linewidth=2.0, zorder=4)\n",
    "        frame += [self.ax.add_artist(joint)]\n",
    "\n",
    "        # draw the information text\n",
    "        text = \"x = {x:>+4.1f} [m], theta = {theta:>+6.1f} [deg], input force = {force:>+6.2f} [N]\".format(x=x, theta=np.rad2deg(theta), force=force)\n",
    "        frame += [self.ax.text(0.5, 1.02, text, ha='center', transform=self.ax.transAxes, fontsize=14, fontfamily='monospace')]\n",
    "\n",
    "        # draw the arrow if the force is not too small\n",
    "        if abs(force) > 1.0e-3:\n",
    "            frame += [self.ax.arrow(cart_x + np.sign(force)*self.cart_w*0.55, cart_y, self.max_length_of_force_arrow*force/self.max_force_abs, 0.0, zorder=3, width=0.1, head_width=0.3, head_length=0.3, fc='black', ec='black')]\n",
    "        self.frames.append(frame)\n",
    "\n",
    "    # rotate shape and return location on the x-y plane.\n",
    "    def _affine_transform(self, xlist: list, ylist: list, angle: float, translation: list=[0.0, 0.0]) -> Tuple[list, list]:\n",
    "        transformed_x = []\n",
    "        transformed_y = []\n",
    "        if len(xlist) != len(ylist):\n",
    "            print(\"[ERROR] xlist and ylist must have the same size.\")\n",
    "            raise AttributeError\n",
    "\n",
    "        for i, xval in enumerate(xlist):\n",
    "            transformed_x.append((xlist[i])*np.cos(angle)-(ylist[i])*np.sin(angle)+translation[0])\n",
    "            transformed_y.append((xlist[i])*np.sin(angle)+(ylist[i])*np.cos(angle)+translation[1])\n",
    "        transformed_x.append(transformed_x[0])\n",
    "        transformed_y.append(transformed_y[0])\n",
    "        return transformed_x, transformed_y\n",
    "\n",
    "    def show_animation(self, interval_ms) -> None:\n",
    "        \"\"\"show animation of the recorded frames\"\"\"\n",
    "        ani = ArtistAnimation(self.fig, self.frames, interval=interval_ms)\n",
    "        html = display.HTML(ani.to_jshtml())\n",
    "        display.display(html)\n",
    "        plt.close()\n",
    "\n",
    "    def save_animation(self, filename, interval, movie_writer=\"ffmpeg\") -> None:\n",
    "        \"\"\"save animation of the recorded frames (ffmpeg required)\"\"\"\n",
    "        ani = ArtistAnimation(self.fig, self.frames, interval=interval, blit=True)\n",
    "        ani.save(filename, writer=movie_writer)\n",
    "\n",
    "## [test simulation] ##\n",
    "sim_step = 100\n",
    "delta_t = 0.02\n",
    "cartpole = CartPole(mass_of_cart=1.0, mass_of_pole=0.1, max_force_abs=10.0)\n",
    "for i in range(sim_step):\n",
    "    cartpole.update(u=[10*np.sin(i/10)], delta_t=delta_t) # u is the control input to the cartpole, [ force[N] ]\n",
    "cartpole.show_animation(interval_ms=delta_t*1000) # show animation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Controller : MPPI Controller\n",
    "\n",
    "### Note\n",
    "The following MPPI implementation follows Algorithms 1 and 2 of the reference paper. \n",
    "\n",
    "### Reference\n",
    "1. G. Williams et al. \"Information-Theoretic Model Predictive Control: Theory and Applications to Autonomous Driving\" \n",
    "    - URL : https://ieeexplore.ieee.org/document/8558663\n",
    "    - PDF : https://arxiv.org/pdf/1707.02342.pdf\n",
    "\n",
    "### Brief overview of MPPI algorithm\n",
    "Here is a general process flow to calculate optimal input with mppi algorithm.\n",
    "\n",
    "**[Step 1]** ramdomly sample input sequence\n",
    "\n",
    "Mean input sequence $U$ and ramdomly sampled input sequence $V$ are defined as follows.  \n",
    "Usually, optimal input sequence on the previous step is used as $U$. \n",
    "$$\n",
    "    \\begin{align}\n",
    "        & (\\mathbf{u}_0, \\mathbf{u}_1, ... \\mathbf{u}_{T-1}) = U \\in \\mathbb{R}^{m \\times T}, \\nonumber \\\\\n",
    "        & (\\mathbf{v}_0, \\mathbf{v}_1, ... \\mathbf{v}_{T-1}) = V \\in \\mathbb{R}^{m \\times T}, \\nonumber \\\\\n",
    "        & \\mathbf{v}_t = \\mathbf{u}_t + \\epsilon_t, \\nonumber \\\\\n",
    "        & \\epsilon_t \\sim \\mathcal{N}(0, \\Sigma).\\nonumber \n",
    "    \\end{align}\n",
    "$$\n",
    "\n",
    "\n",
    "**[Step 2]** predict future states and evaluate cost for each sample\n",
    "\n",
    "We assume a discrete time, continuous state-action dynamical system as a control target.  \n",
    "$\\mathbf{x}$ is a system state, and $\\mathbf{v}$ is a sampled control input.\n",
    "$$\n",
    "\\begin{align}\n",
    "\\mathbf{x}_t  &\\in \\mathbb{R}^{n}, \\nonumber \\\\\n",
    "\\mathbf{x}_{t+1} &= \\mathbf{F}(\\mathbf{x}_t, \\mathbf{v}_t).\\nonumber \n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Then costs (i.e. penalties to be minimized) for sampled sequences $S(V; \\mathbf{x}_0)$ can be evaluated with following formulations.\n",
    "$$\n",
    "    \\begin{align}\n",
    "        & S(V; \\mathbf{x}_0) = C(\\mathcal{H}(V; \\mathbf{x}_0)), \\nonumber \\\\\n",
    "        & C(\\mathbf{x}_0, \\mathbf{x}_1, ... \\mathbf{x}_T) = \\phi(\\mathbf{x}_T) + \\sum_{t=0}^{T-1}c(\\mathbf{x}_t), \\nonumber \\\\\n",
    "        & \\mathcal{H}(V; \\mathbf{x}_0) = \\left( \\mathbf{x}_0, \\mathbf{F}(\\mathbf{x}_0, \\mathbf{v}_0), \\mathbf{F}(\\mathbf{F}(\\mathbf{x}_0, \\mathbf{v}_0), \\mathbf{v}_1), ... \\right).\\nonumber \n",
    "    \\end{align}\n",
    "$$\n",
    "\n",
    "**[Step 3]** calculate weight for each sample sequence\n",
    "\n",
    "Weight for a each sample sequence is derived on the basis of information theory.  \n",
    "There are K sample sequences in total, represented with an index k.  \n",
    "Good control sequence with small cost value get more weight, and vice versa.  \n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "& w(V) = \\frac{1}{\\eta} \\exp\n",
    "\\left( \n",
    "    -\\frac{1}{\\lambda}\n",
    "    \\left(\n",
    "        S(V) + \\lambda(1-\\alpha) \\sum^{T-1}_{t=0} \\mathbf{u}_t^T \\Sigma^{-1} (\\mathbf{u}_t + \\epsilon_t) - \\rho\n",
    "    \\right)\n",
    "\\right) \\nonumber \\\\\n",
    "& \\eta = \n",
    "\\sum_{k=1}^K \\exp\n",
    "\\left( \n",
    "    -\\frac{1}{\\lambda}\n",
    "    \\left(\n",
    "        S(U + \\mathcal{E}_k) + \\lambda(1-\\alpha) \\sum^{T-1}_{t=0} \\mathbf{u}_t^T \\Sigma^{-1} (\\mathbf{u}_t + \\epsilon_t^k) - \\rho\n",
    "    \\right)\n",
    "\\right)\\nonumber \\\\\n",
    "& \\rho = \n",
    "\\min_k \n",
    "\\left( S(V_k) + \\lambda(1-\\alpha) \\sum^{T-1}_{t=0} \\mathbf{u}_t^T \\Sigma^{-1} (\\mathbf{u}_t + \\epsilon_t^k) \\right)\\nonumber\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Note that $\\rho$ is inserted into the formulation to avoid overflow errors during implementation.\n",
    "\n",
    "**[Step 4]** get optimal control input sequence\n",
    "\n",
    "Finally, optimal input trajectory for the next ($i+1$) step is given adding weighted sample sequences to the previous solution.\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "    \\mathbf{u}_t^{i+1} % &= \\mathbb{E}_{\\mathbb{Q}_{\\hat{U}, \\Sigma}}[w(V)\\mathbf{v}_t]\n",
    "                 = u_t^i + \\sum_{k=1}^K w(V_k) \\epsilon_t^k \\nonumber \n",
    "\\end{align}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MPPIControllerForCartPole():\n",
    "    def __init__(\n",
    "            self,\n",
    "            delta_t: float = 0.02,\n",
    "            mass_of_cart: float = 1.0,\n",
    "            mass_of_pole: float = 0.01,\n",
    "            length_of_pole: float = 2.0,\n",
    "            max_force_abs: float = 100.0,\n",
    "            horizon_step_T: int = 100,\n",
    "            number_of_samples_K: int = 1000,\n",
    "            param_exploration: float = 0.0,\n",
    "            param_lambda: float = 50.0,\n",
    "            param_alpha: float = 1.0,\n",
    "            sigma: float = 10.0,\n",
    "            stage_cost_weight: np.ndarray = np.array([5.0, 10.0, 0.1, 0.1]), # weight for [x, theta, x_dot, theta_dot]\n",
    "            terminal_cost_weight: np.ndarray = np.array([5.0, 10.0, 0.1, 0.1]) # weight for [x, theta, x_dot, theta_dot]\n",
    "    ) -> None:\n",
    "        \"\"\"initialize mppi controller for cartpole\"\"\"\n",
    "        # mppi parameters\n",
    "        self.dim_u = 1 # dimension of control input vector\n",
    "        self.T = horizon_step_T # prediction horizon\n",
    "        self.K = number_of_samples_K # number of sample trajectories\n",
    "        self.param_exploration = param_exploration  # constant parameter of mppi\n",
    "        self.param_lambda = param_lambda  # constant parameter of mppi\n",
    "        self.param_alpha = param_alpha # constant parameter of mppi\n",
    "        self.param_gamma = self.param_lambda * (1.0 - (self.param_alpha))  # constant parameter of mppi\n",
    "        self.Sigma = sigma # deviation of noise\n",
    "        self.stage_cost_weight = stage_cost_weight\n",
    "        self.terminal_cost_weight = terminal_cost_weight\n",
    "\n",
    "        # cartpole parameters\n",
    "        self.g = 9.81\n",
    "        self.delta_t = delta_t\n",
    "        self.mass_of_cart = mass_of_cart\n",
    "        self.mass_of_pole = mass_of_pole\n",
    "        self.length_of_pole = length_of_pole\n",
    "        self.max_force_abs = max_force_abs\n",
    "\n",
    "        # mppi variables\n",
    "        self.u_prev = np.zeros((self.T))\n",
    "\n",
    "    def calc_control_input(self, observed_x: np.ndarray) -> Tuple[float, np.ndarray]:\n",
    "        \"\"\"calculate optimal control input\"\"\"\n",
    "        # load privious control input sequence\n",
    "        u = self.u_prev\n",
    "\n",
    "        # set initial x value from observation\n",
    "        x0 = observed_x\n",
    "\n",
    "        # prepare buffer\n",
    "        S = np.zeros((self.K)) # state cost list\n",
    "\n",
    "        # sample noise\n",
    "        epsilon = self._calc_epsilon(self.Sigma, self.K, self.T) # size is self.K x self.T\n",
    "\n",
    "        # loop for 0 ~ K-1 samples\n",
    "        for k in range(self.K):         \n",
    "            # prepare buffer\n",
    "            v = np.zeros((self.T)) # control input sequence with noise\n",
    "\n",
    "            # set initial(t=0) state x i.e. observed state of the cartpole\n",
    "            x = x0\n",
    "\n",
    "            # loop for time step t = 1 ~ T\n",
    "            for t in range(1, self.T+1):\n",
    "\n",
    "                # get control input with noise\n",
    "                if k < (1.0-self.param_exploration)*self.K:\n",
    "                    v[t-1] = u[t-1] + epsilon[k, t-1] # sampling for exploitation\n",
    "                else:\n",
    "                    v[t-1] = epsilon[k, t-1] # sampling for exploration\n",
    "\n",
    "                # update x\n",
    "                x = self._F(x, self._g(v[t-1]))\n",
    "\n",
    "                # add stage cost\n",
    "                S[k] += self._c(x) + self.param_gamma * u[t-1] * (1.0/self.Sigma) * v[t-1]\n",
    "\n",
    "            # add terminal cost\n",
    "            S[k] += self._phi(x)\n",
    "\n",
    "        # compute information theoretic weights for each sample\n",
    "        w = self._compute_weights(S)\n",
    "\n",
    "        # calculate w_k * epsilon_k\n",
    "        w_epsilon = np.zeros((self.T))\n",
    "        for t in range(0, self.T): # loop for time step t = 0 ~ T-1\n",
    "            for k in range(self.K):\n",
    "                w_epsilon[t] += w[k] * epsilon[k, t]\n",
    "\n",
    "        # apply moving average filter for smoothing input sequence\n",
    "        w_epsilon = self._moving_average_filter(xx=w_epsilon, window_size=10)\n",
    "\n",
    "        # update control input sequence\n",
    "        u += w_epsilon\n",
    "\n",
    "        # update privious control input sequence (shift 1 step to the left)\n",
    "        self.u_prev[:-1] = u[1:]\n",
    "        self.u_prev[-1] = u[-1]\n",
    "\n",
    "        # return optimal control input and input sequence\n",
    "        return u[0], u \n",
    "\n",
    "    def _calc_epsilon(self, sigma: float, size_sample: int, size_time_step: int) -> np.ndarray:\n",
    "        \"\"\"sample epsilon\"\"\"\n",
    "        epsilon = np.random.normal(0.0, sigma, (self.K, self.T)) # size is self.K x self.T\n",
    "        return epsilon\n",
    "\n",
    "    def _g(self, v: np.ndarray) -> float:\n",
    "        \"\"\"clamp input\"\"\"\n",
    "        v = np.clip(v, -self.max_force_abs, self.max_force_abs)\n",
    "        return v\n",
    "\n",
    "    def _c(self, x_t: np.ndarray) -> float:\n",
    "        \"\"\"calculate stage cost\"\"\"\n",
    "        # parse x_t\n",
    "        x, x_dot = x_t[0], x_t[2]\n",
    "        theta, theta_dot = x_t[1], x_t[3]\n",
    "        theta = ((theta + np.pi) % (2 * np.pi)) - np.pi # normalize theta to [-pi, pi]\n",
    "\n",
    "        # calculate stage cost # (np.cos(theta)+1.0)\n",
    "        stage_cost = self.stage_cost_weight[0]*x**2 + self.stage_cost_weight[1]*theta**2 + self.stage_cost_weight[2]*x_dot**2 + self.stage_cost_weight[3]*theta_dot**2\n",
    "        return stage_cost\n",
    "\n",
    "    def _phi(self, x_T: np.ndarray) -> float:\n",
    "        \"\"\"calculate terminal cost\"\"\"\n",
    "        # parse x_T\n",
    "        x, x_dot = x_T[0], x_T[2]\n",
    "        theta, theta_dot = x_T[1], x_T[3]\n",
    "        theta = ((theta + np.pi) % (2 * np.pi)) - np.pi # normalize theta to [-pi, pi]\n",
    "\n",
    "        # calculate terminal cost # (np.cos(theta)+1.0)\n",
    "        terminal_cost = self.terminal_cost_weight[0]*x**2 + self.terminal_cost_weight[1]*theta**2 + self.terminal_cost_weight[2]*x_dot**2 + self.terminal_cost_weight[3]*theta_dot**2\n",
    "        return terminal_cost\n",
    "\n",
    "    def _F(self, x_t: np.ndarray, v_t: np.ndarray) -> np.ndarray:\n",
    "        \"\"\"calculate next state of the cartpole\"\"\"\n",
    "        # get previous state variables\n",
    "        x, theta, x_dot, theta_dot = x_t[0], x_t[1], x_t[2], x_t[3]\n",
    "\n",
    "        # prepare params\n",
    "        g = self.g\n",
    "        M = self.mass_of_cart\n",
    "        m = self.mass_of_pole\n",
    "        l = self.length_of_pole\n",
    "        f = v_t\n",
    "        dt = self.delta_t\n",
    "\n",
    "        # get acc. values\n",
    "        temp = (\n",
    "            f + (m*l) * theta_dot**2 * np.sin(theta)\n",
    "        ) / (M + m)\n",
    "        new_theta_ddot = (g * np.sin(theta) - np.cos(theta) * temp) / (\n",
    "            l * (4.0 / 3.0 - m * np.cos(theta)**2 / (M + m))\n",
    "        )\n",
    "        new_x_ddot = temp - (m*l)  * new_theta_ddot * np.cos(theta) / (M+m)\n",
    "\n",
    "        # update pos. values\n",
    "        theta = theta + theta_dot * dt\n",
    "        x = x + x_dot * dt\n",
    "\n",
    "        # update vel. values\n",
    "        theta_dot = theta_dot + new_theta_ddot * dt\n",
    "        x_dot = x_dot + new_x_ddot * dt\n",
    "\n",
    "        # return updated state\n",
    "        x_t_plus_1 = np.array([x, theta, x_dot, theta_dot])\n",
    "        return x_t_plus_1\n",
    "\n",
    "    def _compute_weights(self, S: np.ndarray) -> np.ndarray:\n",
    "        \"\"\"compute weights for each sample\"\"\"\n",
    "        # prepare buffer\n",
    "        w = np.zeros((self.K))\n",
    "\n",
    "        # calculate rho\n",
    "        rho = S.min()\n",
    "\n",
    "        # calculate eta\n",
    "        eta = 0.0\n",
    "        for k in range(self.K):\n",
    "            eta += np.exp( (-1.0/self.param_lambda) * (S[k]-rho) )\n",
    "\n",
    "        # calculate weight\n",
    "        for k in range(self.K):\n",
    "            w[k] = (1.0 / eta) * np.exp( (-1.0/self.param_lambda) * (S[k]-rho) )\n",
    "        return w\n",
    "\n",
    "    def _moving_average_filter(self, xx: np.ndarray, window_size: int) -> np.ndarray:\n",
    "        \"\"\"apply moving average filter for smoothing input sequence\n",
    "        Ref. https://zenn.dev/bluepost/articles/1b7b580ab54e95\n",
    "        Note: The original MPPI paper uses the Savitzky-Golay Filter for smoothing control inputs.\n",
    "        \"\"\"\n",
    "        b = np.ones(window_size)/window_size\n",
    "        xx_mean = np.convolve(xx, b, mode=\"same\")\n",
    "        n_conv = math.ceil(window_size/2)\n",
    "        xx_mean[0] *= window_size/n_conv\n",
    "        for i in range(1, n_conv):\n",
    "            xx_mean[i] *= window_size/(i+n_conv)\n",
    "            xx_mean[-i] *= window_size/(i + n_conv - (window_size % 2)) \n",
    "        return xx_mean\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simulation settings\n",
    "delta_t = 0.02 # [sec]\n",
    "sim_steps = 200 # [steps]\n",
    "print(f\"[INFO] delta_t : {delta_t:.2f}[s] , sim_steps : {sim_steps}[steps], total_sim_time : {delta_t*sim_steps:.2f}[s]\")\n",
    "\n",
    "# initialize a cartpole as a control target\n",
    "cartpole = CartPole(\n",
    "    mass_of_cart = 1.0,\n",
    "    mass_of_pole = 0.01,\n",
    "    length_of_pole = 2.0, \n",
    "    max_force_abs = 100.0,\n",
    ")\n",
    "cartpole.reset(\n",
    "    init_state = np.array([0.0, np.pi, 0.0, 0.0]), # [x[m], theta[rad], x_dot[m/s], theta_dot[rad/s]]\n",
    ")\n",
    "\n",
    "# initialize a mppi controller for the cartpole\n",
    "mppi = MPPIControllerForCartPole(\n",
    "    delta_t = delta_t,\n",
    "    mass_of_cart = 1.0,\n",
    "    mass_of_pole = 0.01,\n",
    "    length_of_pole = 2.0,\n",
    "    max_force_abs = 100.0,\n",
    "    horizon_step_T = 100,\n",
    "    number_of_samples_K = 1000,\n",
    "    param_exploration = 0.0,\n",
    "    param_lambda = 50.0,\n",
    "    param_alpha = 1.0,\n",
    "    sigma = 10.0,\n",
    "    stage_cost_weight    = np.array([5.0, 10.0, 0.1, 0.1]), # weight for [x, theta, x_dot, theta_dot]\n",
    "    terminal_cost_weight = np.array([5.0, 10.0, 0.1, 0.1]), # weight for [x, theta, x_dot, theta_dot]\n",
    ")\n",
    "\n",
    "# simulation loop\n",
    "for i in range(sim_steps):\n",
    "\n",
    "    # get current state of cartpole\n",
    "    current_state = cartpole.get_state()\n",
    "\n",
    "    # calculate input force with MPPI\n",
    "    input_force, input_force_sequence = mppi.calc_control_input(\n",
    "        observed_x = current_state\n",
    "    )\n",
    "\n",
    "    # print current state and input force\n",
    "    print(f\"Time: {i*delta_t:>2.2f}[s], x={current_state[0]:>+3.3f}[m], theta={current_state[1]:>+3.3f}[rad], input force={input_force:>+6.2f}[N]\")\n",
    "\n",
    "    # update states of cartpole\n",
    "    cartpole.update(u=[input_force], delta_t=delta_t)\n",
    "\n",
    "# show animation\n",
    "cartpole.show_animation(interval_ms=int(delta_t * 1000))\n",
    "# save animation\n",
    "cartpole.save_animation(\"mppi_cartpole.mp4\", interval=int(delta_t * 1000), movie_writer=\"ffmpeg\") # ffmpeg is required to write mp4 file"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
